import json
import logging
import asyncio
import aiohttp
from typing import Dict, Any, Optional, List, AsyncGenerator, Tuple
from contextlib import suppress

logger = logging.getLogger(__name__)

class McpClient:
    """
    MCP客户端，用于与MCP服务交互，支持SSE和HTTP两种方式
    """
    
    def __init__(self, server_url: str):
        """
        初始化MCP客户端
        
        Args:
            server_url: MCP服务器URL
        """
        self.server_url = server_url
        self.session = None
        self.streams_context = None
        self.session_context = None
        self.is_connected = False
        self.is_closing = False
        logger.info(f"初始化MCP客户端，服务URL: {server_url}")
    
    @staticmethod
    def extract_server_url(mcp_config: Dict[str, Any]) -> Optional[str]:
        """
        从MCP配置中提取服务器URL
        
        Args:
            mcp_config: MCP配置字典
            
        Returns:
            Optional[str]: 服务器URL，如果未找到则返回None
        """
        if not mcp_config:
            return None
            
        # 尝试从不同位置获取URL
        if isinstance(mcp_config, dict):
            # 处理mcpServers格式
            if 'mcpServers' in mcp_config and isinstance(mcp_config['mcpServers'], dict):
                # 遍历所有服务器配置
                for server_name, server_config in mcp_config['mcpServers'].items():
                    if isinstance(server_config, dict) and 'url' in server_config:
                        logger.info(f"从mcpServers中找到服务器URL: {server_config['url']}")
                        return server_config['url']
            
            # 直接是URL
            if 'url' in mcp_config:
                return mcp_config['url']
            # 在server字段中
            elif 'server' in mcp_config and isinstance(mcp_config['server'], dict):
                if 'url' in mcp_config['server']:
                    return mcp_config['server']['url']
            # 在config字段中
            elif 'config' in mcp_config and isinstance(mcp_config['config'], dict):
                if 'url' in mcp_config['config']:
                    return mcp_config['config']['url']
                elif 'server' in mcp_config['config'] and isinstance(mcp_config['config']['server'], dict):
                    if 'url' in mcp_config['config']['server']:
                        return mcp_config['config']['server']['url']
        
        logger.warning(f"无法从MCP配置中提取服务器URL: {mcp_config}")
        return None
    
    @staticmethod
    async def execute_tool(server_url: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行MCP工具（静态方法）
        
        Args:
            server_url: MCP服务器URL
            parameters: 工具参数，必须包含name字段指定要执行的工具名称
            
        Returns:
            Dict[str, Any]: 执行结果
        """
        logger.info(f"开始执行MCP工具，服务URL: {server_url}")
        logger.info(f"请求参数: {json.dumps(parameters, ensure_ascii=False)}")
        
        # 验证参数
        if not isinstance(parameters, dict):
            error_msg = "参数必须是字典类型"
            logger.error(error_msg)
            return {"error": error_msg}
            
        tool_name = parameters.get("name")
        if not tool_name:
            error_msg = "参数中必须包含name字段指定要执行的工具名称"
            logger.error(error_msg)
            return {"error": error_msg}
            
        logger.info(f"执行工具: {tool_name}")
        
        client = McpClient(server_url)
        if not await client.connect():
            error_msg = "无法连接到MCP服务"
            logger.error(error_msg)
            return {"error": error_msg}
            
        try:
            result = await client.execute(parameters)
            logger.info(f"工具执行结果: {json.dumps(result, ensure_ascii=False)}")
            return result
        except Exception as e:
            error_msg = f"执行工具时发生错误: {str(e)}"
            logger.error(error_msg)
            logger.exception("详细错误栈:")
            return {"error": error_msg}
        finally:
            await client.close()
    
    @staticmethod
    async def get_tool_params(server_url: str) -> Dict[str, Any]:
        """
        获取MCP工具参数（静态方法）
        
        Args:
            server_url: MCP服务器URL
            
        Returns:
            Dict[str, Any]: 工具参数
        """
        client = McpClient(server_url)
        if not await client.connect():
            return {"error": "无法连接到MCP服务"}
            
        try:
            result = await client.get_parameters()
            return result
        finally:
            await client.close()
    
    async def connect(self) -> bool:
        """
        连接到MCP服务器
        
        Returns:
            bool: 连接是否成功
        """
        if self.is_closing:
            logger.warning("正在关闭连接，不允许重新连接")
            return False
            
        if self.is_connected and self.session:
            try:
                if not hasattr(self.session, "closed") or not self.session.closed:
                    logger.info("已有活跃会话，复用现有连接")
                    return True
                else:
                    logger.info("现有会话已标记为关闭，将重新连接")
            except Exception as e:
                logger.warning(f"检查现有会话状态出错: {e}，将重新连接")
        
        try:
            logger.info(f"正在连接MCP服务器: {self.server_url}")
            
            # 使用MCP官方客户端
            from mcp import ClientSession
            from mcp.client.sse import sse_client
            
            # 创建SSE客户端
            self.streams_context = sse_client(url=self.server_url)
            streams = await self.streams_context.__aenter__()
            
            # 创建会话
            self.session_context = ClientSession(*streams)
            self.session = await self.session_context.__aenter__()
            
            # 初始化
            await self.session.initialize()
            logger.info("MCP客户端初始化完成")
            
            self.is_connected = True
            return True
            
        except Exception as e:
            logger.error(f"连接MCP服务器失败: {str(e)}")
            await self.close()
            return False
    
    async def close(self):
        """
        关闭MCP连接
        """
        self.is_closing = True
        
        try:
            logger.info("开始清理MCP连接...")
            
            # 关闭会话
            if self.session_context:
                logger.info("关闭会话...")
                with suppress(Exception):
                    await self.session_context.__aexit__(None, None, None)
                self.session_context = None
                logger.info("会话已关闭")
            
            # 设置会话为None
            self.session = None
            
            # 关闭流
            if self.streams_context:
                logger.info("关闭流...")
                with suppress(Exception):
                    await self.streams_context.__aexit__(None, None, None)
                self.streams_context = None
                logger.info("流已关闭")
            
            logger.info("MCP连接已清理")
        except Exception as e:
            logger.error(f"清理MCP连接时发生错误: {str(e)}")
            logger.exception("详细错误栈:")
        finally:
            self.is_connected = False
            self.is_closing = False
    
    async def check_connection_health(self) -> bool:
        """
        检查与MCP服务器的连接健康状态
        
        Returns:
            bool: 连接是否健康
        """
        try:
            if not self.session:
                logger.warning("会话不存在")
                return False
                
            # 会话闭合检查
            if hasattr(self.session, "closed") and self.session.closed:
                logger.warning("会话已标记为关闭")
                return False
                
            # 实际流可用性检查
            stream_available = True
            if hasattr(self.session, "_write_stream"):
                if hasattr(self.session._write_stream, "closed") and self.session._write_stream.closed:
                    logger.warning("会话写入流已明确标记为关闭")
                    stream_available = False
            
            if stream_available:
                logger.debug("MCP连接健康检查通过")
            else:
                logger.warning("MCP连接健康检查失败")
                
            return stream_available
            
        except Exception as e:
            logger.error(f"检查MCP连接健康时出错: {e}")
            return False
    
    async def list_tools(self) -> Dict[str, Any]:
        """
        获取MCP工具列表
        
        Returns:
            Dict[str, Any]: 工具列表信息
        """
        logger.info(f"开始获取MCP工具列表")
        
        if not self.is_connected or not self.session:
            logger.warning("客户端未连接，尝试连接")
            if not await self.connect():
                return {"error": "无法连接到MCP服务"}
        
        max_retries = 3
        retry_delay = 1
        
        for attempt in range(max_retries):
            try:
                logger.info(f"尝试获取工具列表 (尝试 {attempt + 1}/{max_retries})")
                
                # 检查连接健康状态
                if not await self.check_connection_health():
                    logger.warning("连接不健康，尝试重新连接")
                    if not await self.connect():
                        if attempt < max_retries - 1:
                            logger.info(f"等待 {retry_delay} 秒后重试...")
                            await asyncio.sleep(retry_delay)
                            continue
                        return {"error": "无法连接到MCP服务"}
                
                # 获取工具列表
                logger.info("正在获取工具列表...")
                response = await self.session.list_tools()
                
                if not response:
                    logger.warning("获取工具列表为空")
                    if attempt < max_retries - 1:
                        logger.info(f"等待 {retry_delay} 秒后重试...")
                        await asyncio.sleep(retry_delay)
                        continue
                    return {"error": "获取工具列表为空"}
                
                # 处理工具列表
                tools_list = []
                if hasattr(response, 'tools'):
                    # 使用tools属性
                    for tool in response.tools:
                        # 获取工具参数
                        tool_params = {}
                        if hasattr(tool, 'parameters'):
                            tool_params = tool.parameters
                        elif hasattr(tool, 'inputSchema'):
                            tool_params = tool.inputSchema
                        elif hasattr(tool, 'schema'):
                            tool_params = tool.schema
                        
                        # 如果没有参数，尝试获取默认参数
                        if not tool_params:
                            tool_params = self._get_default_parameters()
                        
                        tools_list.append({
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool_params,
                            "inputSchema": tool_params  # 为了兼容性，同时提供inputSchema
                        })
                        logger.info(f"从响应中找到工具: {tool.name}，参数: {json.dumps(tool_params, ensure_ascii=False)}")
                elif isinstance(response, list):
                    # 直接是工具列表
                    for tool in response:
                        if isinstance(tool, dict) and 'name' in tool:
                            # 获取工具参数
                            tool_params = tool.get("parameters", {}) or tool.get("inputSchema", {}) or tool.get("schema", {})
                            
                            # 如果没有参数，尝试获取默认参数
                            if not tool_params:
                                tool_params = self._get_default_parameters()
                            
                            tools_list.append({
                                "name": tool.get("name", "未命名工具"),
                                "description": tool.get("description", ""),
                                "parameters": tool_params,
                                "inputSchema": tool_params  # 为了兼容性，同时提供inputSchema
                            })
                            logger.info(f"从响应中找到工具: {tool.get('name', '未命名')}，参数: {json.dumps(tool_params, ensure_ascii=False)}")
                elif isinstance(response, dict):
                    # 可能是单个工具或包含工具列表的对象
                    if 'name' in response:
                        # 单个工具
                        # 获取工具参数
                        tool_params = response.get("parameters", {}) or response.get("inputSchema", {}) or response.get("schema", {})
                        
                        # 如果没有参数，尝试获取默认参数
                        if not tool_params:
                            tool_params = self._get_default_parameters()
                        
                        tools_list.append({
                            "name": response.get("name", "未命名工具"),
                            "description": response.get("description", ""),
                            "parameters": tool_params,
                            "inputSchema": tool_params  # 为了兼容性，同时提供inputSchema
                        })
                        logger.info(f"从响应中找到工具: {response.get('name', '未命名')}，参数: {json.dumps(tool_params, ensure_ascii=False)}")
                    elif 'tools' in response and isinstance(response['tools'], list):
                        # 包含工具列表的对象
                        for tool in response['tools']:
                            if isinstance(tool, dict) and 'name' in tool:
                                # 获取工具参数
                                tool_params = tool.get("parameters", {}) or tool.get("inputSchema", {}) or tool.get("schema", {})
                                
                                # 如果没有参数，尝试获取默认参数
                                if not tool_params:
                                    tool_params = self._get_default_parameters()
                                
                                tools_list.append({
                                    "name": tool.get("name", "未命名工具"),
                                    "description": tool.get("description", ""),
                                    "parameters": tool_params,
                                    "inputSchema": tool_params  # 为了兼容性，同时提供inputSchema
                                })
                                logger.info(f"从响应中找到工具: {tool.get('name', '未命名')}，参数: {json.dumps(tool_params, ensure_ascii=False)}")
                
                if tools_list:
                    logger.info(f"成功获取工具列表，共 {len(tools_list)} 个工具")
                    return {"tools": tools_list}
                else:
                    logger.warning("获取工具列表为空")
                    if attempt < max_retries - 1:
                        logger.info(f"等待 {retry_delay} 秒后重试...")
                        await asyncio.sleep(retry_delay)
                        continue
                    return {"error": "获取工具列表为空"}
                    
            except Exception as e:
                logger.error(f"获取工具列表时发生错误: {str(e)}")
                
                # 清理连接
                await self.close()
                
                if attempt < max_retries - 1:
                    logger.info(f"等待 {retry_delay} 秒后重试...")
                    await asyncio.sleep(retry_delay)
                else:
                    logger.error("获取工具列表失败，已达到最大重试次数")
                    return {"error": f"获取工具列表失败: {str(e)}"}
        
        return {"error": "获取工具列表失败，已达到最大重试次数"}
    
    async def execute(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行MCP工具
        
        Args:
            parameters: 工具参数
            
        Returns:
            Dict[str, Any]: 执行结果
        """
        logger.info(f"开始执行MCP工具")
        logger.debug(f"请求参数: {json.dumps(parameters, ensure_ascii=False)}")
        
        if not self.is_connected or not self.session:
            logger.warning("客户端未连接，尝试连接")
            if not await self.connect():
                return {"error": "无法连接到MCP服务"}
        
        max_retries = 3
        retry_delay = 2
        timeout = 30.0  # 超时时间为30秒
        
        for attempt in range(max_retries):
            try:
                logger.info(f"开始执行工具调用 (尝试 {attempt + 1}/{max_retries})")
                
                # 检查连接健康状态
                if not await self.check_connection_health():
                    logger.warning("连接不健康，尝试重新连接")
                    if not await self.connect():
                        if attempt < max_retries - 1:
                            logger.info(f"等待 {retry_delay} 秒后重试...")
                            await asyncio.sleep(retry_delay)
                            continue
                        return {"error": "无法连接到MCP服务"}
                
                # 执行工具调用
                try:
                    tool_name = parameters.get("name")
                    if not tool_name:
                        return {"error": "工具名称不能为空"}
                    
                    logger.info(f"开始调用工具 {tool_name}...")
                    result = await asyncio.wait_for(
                        self.session.call_tool(tool_name, parameters),
                        timeout=timeout
                    )
                    logger.info(f"工具调用成功: {tool_name}")
                    logger.debug(f"工具调用结果: {result}")
                    
                    # 处理CallToolResult类型的结果
                    if hasattr(result, 'content') and isinstance(result.content, list):
                        # 提取文本内容
                        text_content = None
                        for content in result.content:
                            if hasattr(content, 'text'):
                                text_content = content.text
                                break
                        
                        if text_content:
                            try:
                                # 尝试解析JSON字符串
                                parsed_content = json.loads(text_content)
                                return {"result": parsed_content}
                            except json.JSONDecodeError:
                                # 如果不是JSON，直接返回文本
                                return {"result": text_content}
                    
                    # 处理普通字典类型结果
                    if isinstance(result, dict):
                        if "result" not in result and "error" not in result:
                            return {"result": result}
                        return result
                    
                    # 处理其他类型的结果
                    return {"result": str(result)}
                    
                except asyncio.TimeoutError:
                    error_msg = f"工具 {tool_name} 执行超时（{timeout}秒）"
                    logger.error(error_msg)
                    
                    # 检查会话是否因超时而失效
                    if not await self.check_connection_health():
                        logger.warning("因超时导致会话失效，尝试清理并重新连接")
                        await self.close()
                    
                    if attempt < max_retries - 1:
                        logger.info(f"等待 {retry_delay} 秒后重试...")
                        await asyncio.sleep(retry_delay)
                        continue
                    return {"error": error_msg}
                    
            except Exception as e:
                error_msg = f"执行工具调用时发生错误: {str(e)}"
                logger.error(error_msg)
                logger.exception("详细错误栈:")
                
                # 检查是否是连接相关错误
                if "closed" in str(e).lower() or "connection" in str(e).lower() or "stream" in str(e).lower():
                    logger.warning("检测到连接错误，尝试清理并重新连接")
                    await self.close()
                
                if attempt < max_retries - 1:
                    logger.info(f"等待 {retry_delay} 秒后重试...")
                    await asyncio.sleep(retry_delay)
                    continue
                return {"error": error_msg}
        
        return {"error": "执行工具调用失败，已达到最大重试次数"}
    
    async def get_parameters(self) -> Dict[str, Any]:
        """
        获取MCP工具参数
        
        Returns:
            Dict[str, Any]: 工具参数
        """
        logger.info(f"开始获取MCP工具参数")
        
        if not self.is_connected or not self.session:
            logger.warning("客户端未连接，尝试连接")
            if not await self.connect():
                return {"error": "无法连接到MCP服务"}
        
        try:
            # 使用list_tools方法获取工具列表
            result = await self.list_tools()
            
            if "error" in result:
                logger.warning(f"获取工具列表失败: {result['error']}")
                return {"error": result["error"]}
            
            if "tools" not in result or not result["tools"]:
                logger.warning("获取工具列表为空")
                return {"error": "获取工具列表为空"}
            
            # 处理工具列表，提取参数
            parameters = {}
            for tool in result["tools"]:
                if isinstance(tool, dict) and 'name' in tool:
                    tool_name = tool.get("name", "未命名工具")
                    tool_params = tool.get("parameters", {}) or tool.get("inputSchema", {})
                    if tool_params:
                        parameters[tool_name] = tool_params
                        logger.info(f"从工具 {tool_name} 中提取参数")
            
            if parameters:
                logger.info(f"成功获取工具参数，共 {len(parameters)} 个工具")
                return parameters
            else:
                logger.warning("未能获取有效参数，返回默认参数")
                return self._get_default_parameters()
                
        except Exception as e:
            logger.exception(f"获取工具参数时发生错误: {str(e)}")
            return {"error": f"获取工具参数时发生错误: {str(e)}"}
    
    def _get_default_parameters(self) -> Dict[str, Any]:
        """
        获取默认参数
        
        Returns:
            Dict[str, Any]: 默认参数
        """
        logger.info("返回默认参数定义")
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string", 
                    "description": "输入查询文本"
                }
            },
            "required": ["query"]
        } 